#include <iostream>
#include <vector>
#include "NaluDebug.h"
#include "StreamParse.hpp"

static int findStartcode3(unsigned char *buffer)
{
    return (buffer[0]==0 && buffer[1]==0 && buffer[2]==1);
}

static int findStartcode4(unsigned char *buffer)
{
    return (buffer[0]==0 && buffer[1]==0 && buffer[2]==0 && buffer[3]==1);
}

StreamParse::StreamParse(int32_t type)
: m_hH264(NULL), 
  m_hH265(NULL), 
  pix_width(0),
  pix_height(0)
{
    if (type == 0){
        m_hH264 = h264_new();
        m_hH265 = NULL;
    }else if (type == 1){
        m_hH264 = NULL;
        m_hH265 = h265_new();
    }else{
        parse_error("not support type\n");
        return ;
    }
    m_nType = type;
}

StreamParse::Ptr StreamParse::CreateShared(int32_t type)
{
    // struct EnableMakeShared : public StreamParse {
    //     EnableMakeShared(int32_t type) : StreamParse(type) {}
    // };
    // return std::make_shared<EnableMakeShared>(type);
    return Ptr(new StreamParse(type));
}

StreamParse::Ptr StreamParse::CreateH264Shared()
{
    return Ptr(new StreamParse(NALU_VIDEO_H264));
}

StreamParse::Ptr StreamParse::CreateH265Shared()
{
    return Ptr(new StreamParse(NALU_VIDEO_H265));
}

StreamParse::Ptr StreamParse::CreateHEVCShared()
{
    return Ptr(new StreamParse(NALU_VIDEO_HEVC));
}

StreamParse::~StreamParse()
{
    if (m_hH264 != NULL){
        h264_free(m_hH264);
    }
    if (m_hH265 != NULL){
        h265_free(m_hH265);
    }
}

void StreamParse::getVideoInfo(videoinfo_t* videoinfo)
{
    std::lock_guard<std::mutex> lock(mtx);
    if (videoinfo == NULL){
        return ;
    }
    videoinfo->nalu_count = 0;
    videoinfo->file_size = 0;
    if (m_nType == 0){
        memcpy(videoinfo, m_hH264->info, sizeof(videoinfo_t));
        videoinfo->width = pix_width;
        videoinfo->height = pix_height;
        return;
    }else if (m_nType == 1){
        memcpy(videoinfo, m_hH265->info, sizeof(videoinfo_t));
        videoinfo->width = pix_width;
        videoinfo->height = pix_height;
        return;
    }
    memset(videoinfo, 0 ,sizeof(videoinfo_t));
    videoinfo->type = 2;
    parse_error("not support type\n");
}

int32_t StreamParse::width() const
{
    return pix_width;
}

int32_t StreamParse::height() const
{
    return pix_height;
}

float StreamParse::fps() const
{
    if (m_nType == 0) {
        return m_hH264->info->fps;
    }
    else if(m_nType == 1) {
        return m_hH265->info->fps;
    }
    return 0;
}

int32_t StreamParse::streamCodecId()
{
    return m_nType;
}

std::string StreamParse::streamCodecName()
{
    std::lock_guard<std::mutex> lock(mtx);
    std::string codec_name = "unknown";
    if (m_nType == 0) {
        codec_name = "H264";
    }
    else if(m_nType == 1) {
        codec_name = "H265";
    }
    return codec_name;
}

void StreamParse::DebugSPS() 
{
    std::lock_guard<std::mutex> lock(mtx);
    if (m_nType == 1) {
        h265_debug_sps(m_hH265->sps);
    }
    else if(m_nType == 0) {
        h264_debug_sps(m_hH264->sps);
    }
}

void StreamParse::DebugPPS()
{
    std::lock_guard<std::mutex> lock(mtx);
    if (m_nType == 1) {
        h265_debug_pps(m_hH265->pps);
    }
    else if(m_nType == 0) {
        h264_debug_pps(m_hH264->pps);
    }
}

void StreamParse::DebugVPS()
{
    std::lock_guard<std::mutex> lock(mtx);
    if (m_nType == 1) {
        h265_debug_vps(m_hH265->vps);
    }
}

bool StreamParse::parseStream(const void* data, size_t size)
{
    if(data == NULL || size < 5){
        return false;
    }
    int32_t start_len = 0;
    const uint8_t* buffer = static_cast<const uint8_t*>(data);
    if(!_keyframe(m_nType, buffer, size, &start_len) && 
         !_configFrame(m_nType, buffer, size, &start_len))return false;
    std::vector<NALU_t> nalus;
    nalus.clear();

    uint64_t file_position = 0;  // Current position in file
    uint64_t last_nalu_pos = 0;  // Start position of the last NALU found
    size_t bytes_read = size;

    for (size_t i = 0; i < bytes_read; i++) {
        if (i < bytes_read - 4 &&
            buffer[i] == 0x00 && buffer[i + 1] == 0x00 &&
            ((buffer[i + 2] == 0x01) || (buffer[i + 2] == 0x00 && buffer[i + 3] == 0x01))) {
            
            // Complete the previous NALU size calculation before recording new one
            if (nalus.size() > 0) {
                nalus[nalus.size() - 1].len = file_position + i - last_nalu_pos;
            }

            // Record the current NALU header position
            last_nalu_pos = file_position + i;
            NALU_t new_nalu;
            new_nalu.offset = last_nalu_pos;
            int startcodeLen = (buffer[i + 2] == 0x01) ? 2 : 3;
            new_nalu.startcodeLen = startcodeLen + 1;
            new_nalu.type = m_nType;
            new_nalu.num = nalus.size();
            if(m_nType == 0){
                new_nalu.nalType = buffer[i + new_nalu.startcodeLen] & 0x1f;
                new_nalu.first_mb_in_slice = (buffer[i + new_nalu.startcodeLen + 1] & 0xff);
            }else{
                new_nalu.nalType = (buffer[i + new_nalu.startcodeLen] >> 1) & 0x3f;
                new_nalu.first_mb_in_slice = (buffer[i + new_nalu.startcodeLen + 2] & 0xff);
            }
            nalus.push_back(new_nalu);
            i += startcodeLen;
        }
    }
    file_position += bytes_read;
    if (nalus.size() > 0) {
        nalus[nalus.size() - 1].len = file_position - last_nalu_pos;
    }
    std::lock_guard<std::mutex> lock(mtx);
    for(size_t i = 0; i < nalus.size(); i++){
        parseNALU(buffer, size, &nalus[i]);
        if (m_nType == 0){
            if(m_hH264->info->width > 0 && m_hH264->info->height > 0){
                break;
            }
        }else if (m_nType == 1){
            if(m_hH265->info->width > 0 && m_hH265->info->height > 0){
                break;
            }
        }
    }
    return true;
}

bool StreamParse::parseNALU(const uint8_t* data, size_t size, NALU_t* nalu)
{
    if(nalu == NULL)return false;
    if(nalu->len == 0 || (nalu->startcodeLen  != 3 && nalu->startcodeLen != 4)){
        parse_error("nalu->len == 0 || nalu->startcodeLen %d error\n", nalu->startcodeLen);
        return false;
    }
    if(nalu->offset + nalu->len > size){
        parse_error("nalu->offset + nalu->len > size\n");
        return false;
    }
    // parse_info("nalutype: %d %d %d \n", nalu->nalType, configFrame(nalu) ,keyframe(nalu));
    if(!configFrame(nalu) && !keyframe(nalu))return 0;
    uint8_t *buffer = (uint8_t *)malloc(sizeof(uint8_t) * nalu->len);
    if(buffer == NULL){
        parse_error("malloc error\n");
        return false;
    }
    memcpy(buffer, data + nalu->offset, nalu->len);
    if(m_nType == 0){
        m_hH264->sh->read_slice_type = 1;
        read_nal_unit(m_hH264, &buffer[nalu->startcodeLen], nalu->len - nalu->startcodeLen);
        nalu->nalType = m_hH264->nal->nal_unit_type;
        nalu->sliceType = m_hH264->sh->slice_type;
        m_hH264->sh->read_slice_type = 0;
        if(m_hH264->info->width > 0 && m_hH264->info->height > 0){
            pix_width = m_hH264->info->width;
            pix_height = m_hH264->info->height;
        }
    }else{
        m_hH265->sh->read_slice_type = 1;
        h265_read_nal_unit(m_hH265, &buffer[nalu->startcodeLen], nalu->len - nalu->startcodeLen);
        nalu->nalType = m_hH265->nal->nal_unit_type;
        nalu->sliceType = m_hH265->sh->slice_type;
        m_hH265->sh->read_slice_type = 0;
        if(m_hH265->info->width > 0 && m_hH265->info->height > 0){
            pix_width = m_hH265->info->width;
            pix_height = m_hH265->info->height;
        }
    }
    free(buffer);
    buffer = NULL;
    return nalu->len;
}
